#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Date : 2024/10/28
# @Author : 潘世友


from itertools import combinations

"""
主要功能:
1. update_class_attr: 更新class实例的属性变量

使用方法:
from rmtools.utils.code_simplifier import update_class_attr

class MyClass:
    def __init__(self, a, b):
        self.a = a
        self.b = b

    def my_method(self):
        update_class_attr(self, a=3, b=4)
        print(self.a, self.b) # 3 4
        
2. gen_xslist: 生成指标组合列表

使用方法:
from rmtools.utils.code_simplifier import gen_xslist

list_xs = ['a', 'b', 'c', 'd']
gen_xslist(list_xs, 2) # [['a', 'b'], ['a', 'c'], ['a', 'd'], ['b', 'c'], ['b', 'd'], ['c', 'd']]

list_xs = [['a', 'b'], ['c', 'd'], ['e', 'f', 'g']]
gen_xslist(list_xs, 2) # [['a', 'b'], ['c', 'd']]

list_xs = ['a', 'b', 'c', 'd', 'e']
gen_xslist(list_xs, range(2, 4), xs_base=['a', 'b', 'c']) # [['a', 'b'], ['a', 'c'], ['b', 'c'], ['a', 'b', 'c']]
"""


def update_class_attr(instance, **kwargs):
    """ 更新class实例的属性变量 """
    for k, v in kwargs.items():
        if hasattr(instance, k):
            setattr(instance, k, v)


def _gen_xs_simple(xs_: list[str], n_: int | range) -> list[list[str]]:
    if isinstance(n_, int):
        return [list(x) for x in combinations(xs_, n_)]
    elif isinstance(n_, range):
        list_xs = []
        for n_x in n_:
            list_xs += _gen_xs_simple(xs_, n_x)
        return list_xs


def _gen_xs_wdrop(xs_: list[str], n_: int | range, drop_: list[str] | str | None = None) -> list[list[str]]:
    if drop_ is None:
        return _gen_xs_simple(xs_, n_)
    elif isinstance(drop_, list):
        xs_ = [x for x in xs_ if x not in drop_]
        return _gen_xs_simple(xs_, n_)
    elif isinstance(drop_, str):
        xs_ = [x for x in xs_ if x != drop_]
        return _gen_xs_simple(xs_, n_)
    else:
        raise ValueError('xs或xs_drop参数格式错误')


def _gen_xs_wbase(xs_: list[str], n_: int | range, base_: list[str] | None = None) -> list[list[str]]:
    """ xs的组合，仅聚焦存在的指标清单 """
    if base_ is None:
        return _gen_xs_simple(xs_, n_)
    elif isinstance(base_, list):
        xs_ = [x for x in xs_ if x in base_]
        return _gen_xs_simple(xs_, n_)
    else:
        raise ValueError('xs或xs_base参数格式错误')


def _gen_xs_wbd(xs_: list[str],
                n_: int | range,
                base_: list[str] | None = None,
                drop_: list[str] | str | None = None) -> list[list[str]]:
    """ xs的组合，聚焦存在的指标清单，排除xs_drop """
    if base_ is None and drop_ is None:
        return _gen_xs_simple(xs_, n_)
    elif base_ is not None and drop_ is None:
        return _gen_xs_wbase(xs_, n_, base_)
    elif base_ is None and drop_ is not None:
        return _gen_xs_wdrop(xs_, n_, drop_)
    else:
        base_ = [x for x in base_ if x not in drop_]
        return _gen_xs_wbase(xs_, n_, base_)


def _gen_xslist_wbd(xs_: list[list[str]],
                    n_: int | range | None = None,
                    base_: list[str] | None = None,
                    drop_: list[str] | str | None = None) -> list[list[str]]:
    """ xs的组合，聚焦存在的指标清单，排除xs_drop """
    if n_ is None:
        pass
    elif isinstance(n_, int):
        xs_ = [x for x in xs_ if len(x) == n_]
    elif isinstance(n_, range):
        xs_ = [x for x in xs_ if len(x) in n_]
    else:
        raise ValueError('n_xs参数格式错误')

    if isinstance(drop_, str):
        drop_ = [drop_]
    else:
        pass

    if base_ is None and drop_ is None:
        return xs_
    elif base_ is not None and drop_ is None:
        return [x for x in xs_ if set(x) <= set(base_)]
    elif base_ is None and drop_ is not None:
        return [x for x in xs_ if not set(drop_) & set(x)]
    else:
        base_ = [x for x in base_ if x not in drop_]
        return [x for x in xs_ if set(x) <= set(base_)]


def gen_xslist(xs: list[str] | list[list[str]],
               n_xs: int | range | None = None,
               xs_base: list[str] | None = None,
               xs_drop: list[str] | str | None = None) -> list[list[str]]:
    """ xs的组合，聚焦存在的指标清单，排除xs_drop

    Parameters
    ----------
    xs: list[str] | list[list[str]]
        待组合的指标清单，可以是字符串列表或指标组合列表
    n_xs: int | range | None, optional
        组合的长度，默认None，即所有长度的组合
    xs_base: list[str] | None, optional
        聚焦的指标清单，默认None，即所有指标
    xs_drop: list[str] | str | None, optional
        待排除的指标清单，可以是字符串或列表，默认None，即不排除任何指标

    Returns
    -------
    list[list[str]]
        指标组合的列表

    Examples
    --------
    >>> list_xs = ['a', 'b', 'c', 'd']
    >>> gen_xslist(list_xs, 2)
    [['a', 'b'], ['a', 'c'], ['a', 'd'], ['b', 'c'], ['b', 'd'], ['c', 'd']]
    >>> gen_xslist(list_xs, 2, xs_base=['a', 'b', 'c'])
    [['a', 'b'], ['a', 'c'], ['b', 'c']]
    >>> gen_xslist(list_xs, 2, xs_drop='a')
    [['b', 'c'], ['b', 'd'], ['c', 'd']]
    >>> gen_xslist(list_xs, 2, xs_base=['a', 'b'], xs_drop='a')
    [['b', 'c'], ['b', 'd']]
    >>> list_xs = [['a', 'b'], ['c', 'd'], ['e', 'f', 'g']]
    >>> gen_xslist(list_xs, 2)
    [['a', 'b'], ['c', 'd']]
    """

    if all(isinstance(x, str) for x in xs):
        return _gen_xs_wbd(xs, n_xs, xs_base, xs_drop)
    elif all(isinstance(x, list) for x in xs):
        return _gen_xslist_wbd(xs, n_xs, xs_base, xs_drop)
    else:
        raise ValueError('xs参数格式错误')
